import ast
from graphviz import Digraph


class CFGNode:
    """CFG节点类"""

    def __init__(self, node_id, label):
        self.id = node_id
        self.label = label
        self.visited = False


class CFGEdge:
    """CFG边类"""

    def __init__(self, src, dst):
        self.src = src
        self.dst = dst


class CFGBuilder(ast.NodeVisitor):
    """基于AST的CFG构建器"""

    def __init__(self):
        self.nodes = []
        self.edges = []
        self.node_counter = 0
        self.current_nodes = []
        self.node_map = {}  # 存储节点标签到ID的映射

    def new_node(self, label=""):
        """创建新节点"""
        node_id = self.node_counter
        self.node_counter += 1
        self.nodes.append({"id": node_id, "label": label})
        self.node_map[label] = node_id
        return node_id

    def add_edge(self, src, dst):
        """添加边"""
        if src is not None and dst is not None and (src, dst) not in self.edges:
            self.edges.append((src, dst))

    def build_cfg(self, source_code):
        """构建CFG主方法"""
        # 解析源码为AST
        tree = ast.parse(source_code)

        # 创建入口节点
        entry_id = self.new_node("Entry")
        exit_id = self.new_node("Exit")
        self.current_nodes = [entry_id]

        # 遍历AST节点构建CFG
        for stmt in tree.body:
            self.visit(stmt)

        # 添加到出口节点的连接
        for node_id in self.current_nodes:
            self.add_edge(node_id, exit_id)

        return self.generate_adjacency_matrix(), self.nodes, self.edges

    def generate_adjacency_matrix(self):
        """生成邻接矩阵"""
        size = len(self.nodes)
        matrix = [[0] * size for _ in range(size)]
        for src, dst in self.edges:
            matrix[src][dst] = 1
        return matrix

    def visit_Assign(self, node):
        """处理赋值语句"""
        assign_label = f"Assign {ast.unparse(node.targets[0])} = {ast.unparse(node.value)}"
        assign_id = self.new_node(assign_label)

        new_current = []
        for src in self.current_nodes:
            self.add_edge(src, assign_id)
        new_current.append(assign_id)
        self.current_nodes = new_current
        self.generic_visit(node)

    def visit_Expr(self, node):
        """处理表达式"""
        expr_label = f"Expr {ast.unparse(node.value)}"
        expr_id = self.new_node(expr_label)

        new_current = []
        for src in self.current_nodes:
            self.add_edge(src, expr_id)
        new_current.append(expr_id)
        self.current_nodes = new_current

    def visit_If(self, node):
        """处理if条件分支"""
        cond_label = f"If {ast.unparse(node.test)}"
        cond_id = self.new_node(cond_label)

        # 添加从当前节点到条件节点的边
        for src in self.current_nodes:
            self.add_edge(src, cond_id)

        # 处理then分支
        then_id = self.new_node("Then")
        self.add_edge(cond_id, then_id)
        self.current_nodes = [then_id]
        for stmt in node.body:
            self.visit(stmt)

        # 保存then分支结束后的节点
        then_end_nodes = self.current_nodes[:]

        # 处理else分支
        self.current_nodes = [cond_id]
        if node.orelse:
            else_id = self.new_node("Else")
            self.add_edge(cond_id, else_id)
            self.current_nodes = [else_id]
            for stmt in node.orelse:
                self.visit(stmt)

        # 合并分支
        merge_id = self.new_node("Merge")
        # 添加从then分支到合并点的边
        for src in then_end_nodes:
            self.add_edge(src, merge_id)
        # 添加从else分支到合并点的边
        for src in self.current_nodes:
            self.add_edge(src, merge_id)

        self.current_nodes = [merge_id]

    def visit_For(self, node):
        """处理for循环"""
        loop_label = f"For {ast.unparse(node.target)} in {ast.unparse(node.iter)}"
        loop_id = self.new_node(loop_label)

        # 添加从当前节点到循环节点的边
        for src in self.current_nodes:
            self.add_edge(src, loop_id)

        # 循环体起始节点
        body_start_id = self.new_node("Loop Body")
        self.add_edge(loop_id, body_start_id)
        self.current_nodes = [body_start_id]

        # 处理循环体
        for stmt in node.body:
            self.visit(stmt)

        # 循环体回到条件判断
        for src in self.current_nodes:
            self.add_edge(src, loop_id)

        # 循环结束后的新节点
        after_loop_id = self.new_node("After Loop")
        self.add_edge(loop_id, after_loop_id)
        self.current_nodes = [after_loop_id]

    def visit_While(self, node):
        """处理while循环"""
        loop_label = f"While {ast.unparse(node.test)}"
        loop_id = self.new_node(loop_label)

        # 添加从当前节点到循环节点的边
        for src in self.current_nodes:
            self.add_edge(src, loop_id)

        # 循环体起始节点
        body_start_id = self.new_node("Loop Body")
        self.add_edge(loop_id, body_start_id)
        self.current_nodes = [body_start_id]

        # 处理循环体
        for stmt in node.body:
            self.visit(stmt)

        # 循环体回到条件判断
        for src in self.current_nodes:
            self.add_edge(src, loop_id)

        # 循环结束后的新节点
        after_loop_id = self.new_node("After Loop")
        self.add_edge(loop_id, after_loop_id)
        self.current_nodes = [after_loop_id]

    def visit_FunctionDef(self, node):
        """处理函数定义"""
        func_entry_id = self.new_node(f"Function: {node.name}")
        prev_current = self.current_nodes
        self.current_nodes = [func_entry_id]

        for stmt in node.body:
            self.visit(stmt)

        # 保留函数退出节点
        func_exit_id = self.new_node(f"End of {node.name}")
        for src in self.current_nodes:
            self.add_edge(src, func_exit_id)
        self.current_nodes = [func_exit_id]

    def visit_Return(self, node):
        """处理返回语句"""
        return_label = f"Return {ast.unparse(node.value) if node.value else ''}"
        return_id = self.new_node(return_label)

        new_current = []
        for src in self.current_nodes:
            self.add_edge(src, return_id)
        new_current.append(return_id)
        self.current_nodes = new_current

    def save_cfg_as_png(self, filename):
        """保存CFG为PNG格式"""
        dot = self.to_dot()
        try:
            dot.render(filename, format='png', cleanup=True)
        except Exception as e:
            print(f"保存CFG时发生错误: {e}")

    def to_dot(self):
        """生成DOT图形表示"""
        dot = Digraph()
        dot.attr(rankdir='TB')

        # 设置节点样式
        dot.node_attr.update(shape='box', style='filled', fillcolor='white', fontname='Segoe UI')

        # 添加节点
        for node in self.nodes:
            dot.node(str(node["id"]), str(node["label"]),
                     fillcolor='lightblue' if 'If' in str(node["label"]) else 'lightgreen')

        # 添加边
        for src, dst in self.edges:
            dot.edge(str(src), str(dst))

        return dot
